import time
import sys

from unitree_sdk2py.core.channel import ChannelPublisher, ChannelFactoryInitialize
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_
from unitree_sdk2py.utils.crc import CRC
from unitree_sdk2py.utils.thread import RecurrentThread
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient

import numpy as np

kPi = 3.141592654
kPi_2 = 1.57079632

class G1JointIndex:
    # 左腿关节
    LeftHipPitch = 0      # 左髋关节俯仰
    LeftHipRoll = 1       # 左髋关节滚转
    LeftHipYaw = 2        # 左髋关节偏航
    LeftKnee = 3          # 左膝关节
    LeftAnklePitch = 4    # 左踝关节俯仰
    LeftAnkleB = 4        # 左踝关节B（与俯仰相同）
    LeftAnkleRoll = 5     # 左踝关节滚转
    LeftAnkleA = 5        # 左踝关节A（与滚转相同）

    # 右腿关节
    RightHipPitch = 6     # 右髋关节俯仰
    RightHipRoll = 7      # 右髋关节滚转
    RightHipYaw = 8       # 右髋关节偏航
    RightKnee = 9         # 右膝关节
    RightAnklePitch = 10  # 右踝关节俯仰
    RightAnkleB = 10      # 右踝关节B（与俯仰相同）
    RightAnkleRoll = 11   # 右踝关节滚转
    RightAnkleA = 11      # 右踝关节A（与滚转相同）

    # 腰部关节
    WaistYaw = 12         # 腰部偏航
    WaistRoll = 13        # 腰部滚转（对于 g1 23dof/29dof 锁定腰部无效）
    WaistA = 13           # 腰部关节A（对于 g1 23dof/29dof 锁定腰部无效）
    WaistPitch = 14       # 腰部俯仰（对于 g1 23dof/29dof 锁定腰部无效）
    WaistB = 14           # 腰部关节B（对于 g1 23dof/29dof 锁定腰部无效）

    # 左臂关节
    LeftShoulderPitch = 15  # 左肩关节俯仰
    LeftShoulderRoll = 16   # 左肩关节滚转
    LeftShoulderYaw = 17    # 左肩关节偏航
    LeftElbow = 18          # 左肘关节
    LeftWristRoll = 19      # 左腕关节滚转
    LeftWristPitch = 20     # 左腕关节俯仰（对于 g1 23dof 无效）
    LeftWristYaw = 21       # 左腕关节偏航（对于 g1 23dof 无效）

    # 右臂关节
    RightShoulderPitch = 22 # 右肩关节俯仰
    RightShoulderRoll = 23  # 右肩关节滚转
    RightShoulderYaw = 24   # 右肩关节偏航
    RightElbow = 25         # 右肘关节
    RightWristRoll = 26     # 右腕关节滚转
    RightWristPitch = 27    # 右腕关节俯仰（对于 g1 23dof 无效）
    RightWristYaw = 28      # 右腕关节偏航（对于 g1 23dof 无效）

    kNotUsedJoint = 29      # 未使用的关节（权重）

class Custom:
    def __init__(self):
        self.time_ = 0.0
        self.control_dt_ = 0.02  # 控制周期为20ms
        self.duration_ = 3.0     # 每个阶段持续3秒
        self.counter_ = 0
        self.weight = 0.         # 权重初始化
        self.weight_rate = 0.2   # 权重变化速率
        self.kp = 60.            # 比例增益
        self.kd = 1.5            # 微分增益
        self.dq = 0.             # 目标速度
        self.tau_ff = 0.         # 前馈力矩
        self.mode_machine_ = 0   # 机器人模式
        self.low_cmd = unitree_hg_msg_dds__LowCmd_()  # 低级命令
        self.low_state = None    # 低级状态
        self.first_update_low_state = False  # 是否首次更新低级状态
        self.crc = CRC()         # CRC校验
        self.done = False        # 动作是否完成标志

        # 目标关节位置（弧度）
        self.target_pos = [
            0.0,      kPi_2,  0.0,    kPi_2,  0.0,  # 左臂目标位置
            0.0,     -kPi_2,  0.0,    kPi_2,  0.0,  # 右臂目标位置
            0.0,      0.0,    0.0                 # 腰部目标位置
        ]

        # 需要控制的关节索引
        self.arm_joints = [
          G1JointIndex.LeftShoulderPitch,  G1JointIndex.LeftShoulderRoll,
          G1JointIndex.LeftShoulderYaw,    G1JointIndex.LeftElbow,
          G1JointIndex.LeftWristRoll,
          G1JointIndex.RightShoulderPitch, G1JointIndex.RightShoulderRoll,
          G1JointIndex.RightShoulderYaw,   G1JointIndex.RightElbow,
          G1JointIndex.RightWristRoll,
          G1JointIndex.WaistYaw,
          G1JointIndex.WaistRoll,
          G1JointIndex.WaistPitch
        ]

    def Init(self):
        # 创建发布者，用于发送控制命令
        self.arm_sdk_publisher = ChannelPublisher("rt/arm_sdk", LowCmd_)
        self.arm_sdk_publisher.Init()

        # 创建订阅者，用于接收机器人状态
        self.lowstate_subscriber = ChannelSubscriber("rt/lowstate", LowState_)
        self.lowstate_subscriber.Init(self.LowStateHandler, 10)

    def Start(self):
        # 创建线程，周期性调用 LowCmdWrite 方法
        self.lowCmdWriteThreadPtr = RecurrentThread(
            interval=self.control_dt_, target=self.LowCmdWrite, name="control"
        )
        # 等待首次接收到低级状态
        while self.first_update_low_state == False:
            time.sleep(1)

        if self.first_update_low_state == True:
            self.lowCmdWriteThreadPtr.Start()

    def LowStateHandler(self, msg: LowState_):
        # 接收低级状态消息
        self.low_state = msg

        # 标记首次接收到低级状态
        if self.first_update_low_state == False:
            self.first_update_low_state = True
        
    def LowCmdWrite(self):
        # 更新时间
        self.time_ += self.control_dt_

        if self.time_ < self.duration_:
            # [阶段1]: 将机器人设置为零姿态
            self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q = 1  # 启用 arm_sdk
            for i, joint in enumerate(self.arm_joints):
                ratio = np.clip(self.time_ / self.duration_, 0.0, 1.0)
                self.low_cmd.motor_cmd[joint].tau = 0.  # 力矩为0
                self.low_cmd.motor_cmd[joint].q = (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值到零姿态
                self.low_cmd.motor_cmd[joint].dq = 0.  # 速度为0
                self.low_cmd.motor_cmd[joint].kp = self.kp  # 设置比例增益
                self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置微分增益

        elif self.time_ < self.duration_ * 3:
            # [阶段2]: 举起手臂
            for i, joint in enumerate(self.arm_joints):
                ratio = np.clip((self.time_ - self.duration_) / (self.duration_ * 2), 0.0, 1.0)
                self.low_cmd.motor_cmd[joint].tau = 0.  # 力矩为0
                self.low_cmd.motor_cmd[joint].q = ratio * self.target_pos[i] + (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值到目标位置
                self.low_cmd.motor_cmd[joint].dq = 0.  # 速度为0
                self.low_cmd.motor_cmd[joint].kp = self.kp  # 设置比例增益
                self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置微分增益

        elif self.time_ < self.duration_ * 6:
            # [阶段3]: 将机器人恢复到零姿态
            for i, joint in enumerate(self.arm_joints):
                ratio = np.clip((self.time_ - self.duration_ * 3) / (self.duration_ * 3), 0.0, 1.0)
                self.low_cmd.motor_cmd[joint].tau = 0.  # 力矩为0
                self.low_cmd.motor_cmd[joint].q = (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值回到零姿态
                self.low_cmd.motor_cmd[joint].dq = 0.  # 速度为0
                self.low_cmd.motor_cmd[joint].kp = self.kp  # 设置比例增益
                self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置微分增益

        elif self.time_ < self.duration_ * 7:
            # [阶段4]: 释放 arm_sdk
            for i, joint in enumerate(self.arm_joints):
                ratio = np.clip((self.time_ - self.duration_ * 6) / (self.duration_), 0.0, 1.0)
                self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q = (1 - ratio)  # 禁用 arm_sdk

        else:
            # 动作完成
            self.done = True
  
        # 计算 CRC 校验并发送命令
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        self.arm_sdk_publisher.Write(self.low_cmd)

if __name__ == '__main__':
    print("WARNING: Please ensure there are no obstacles around the robot while running this example.")
    input("Press Enter to continue...")

    if len(sys.argv) > 1:
        ChannelFactoryInitialize(0, sys.argv[1])
    else:
        ChannelFactoryInitialize(0)

    custom = Custom()
    custom.Init()
    custom.Start()

    while True:        
        time.sleep(1)
        if custom.done: 
            print("Done!")
            sys.exit(-1)